from typing import Tuple, List, AnyStr

import numpy as np
import pandas as pd
from service.mysql_service import mysql_upload_data, mysql_search_by_id
from config import *
from service.bert_service import sentence_embedding
# from pymilvus_orm import connections, utility
# from pymilvus_orm.mutation import MutationResult
from util.milvus_helpers import MilvusHelper

milvus_client = MilvusHelper(host=MILVUS_HOST, port=MILVUS_PORT, vector_dimension=BERT_DIMENSION)


def milvus_insert(vector: np.ndarray) -> list:
    return milvus_client.insert(collection_name=MILVUS_QA_TABLE, vectors=vector)


def milvus_upload_data():
    """
    上传数据, 会上传数据到milvus和mysql/tidb, 是覆盖式的
    :return: 返回后的数据
    """
    data = pd.read_csv(KNOWLEDGE_FILE_PATH)[['question', 'answer']].dropna(how="any")
    questions = [each for each in data['question'].values]
    vectors = sentence_embedding(questions)
    if milvus_client.has_collection(MILVUS_QA_TABLE):
        milvus_client.delete_collection(MILVUS_QA_TABLE)
    ids = milvus_insert(vectors)
    data['id'] = ids
    print(data)
    mysql_upload_data(data)
    return data


def milvus_similarity_search(question: str) -> Tuple[List[AnyStr], List[float]]:
    """
    根据问题从Milvus库中检索出相似的问题
    :param question:
    :return: 相似问题的id, 相似问题的相似度
    """
    feat = sentence_embedding([question])
    results = milvus_client.search_vectors(MILVUS_QA_TABLE, feat, ANSWER_TOP_K)
    return [str(x.id) for x in results[0]], [float(x.distance) for x in results[0]]


def milvus_answer_search(question: str) -> List[AnyStr]:
    return [mysql_search_by_id(_id) for _id in milvus_similarity_search(question)[0]]


def milvus_delete():
    """
    删除数据
    :return:
    """
    milvus_client.delete_collection(MILVUS_QA_TABLE)


